# -*- coding: utf-8 -*-
"""
@Time ： 2024/1/23 21:52
@Auth ： jingo
@File ：NormalizationBatch.py
@IDE ：PyCharm
@Purpose：栅格数据归一化

"""
import numpy as np
import numpy.ma as ma
from osgeo import gdal


def Get_data(filepath, max_min = False):
    # print("------------------------------------影像基础信息------------------------------------")
    ds = gdal.Open(filepath)  # 打开数据集dataset
    ds_width = ds.RasterXSize  # 获取数据宽度
    ds_height = ds.RasterYSize  # 获取数据高度
    ds_bands = ds.RasterCount  # 获取波段数
    ds_geo = ds.GetGeoTransform()  # 获取仿射地理变换参数
    ds_prj = ds.GetProjection()  # 获取投影信息
    ds_nodata = ds.GetRasterBand(1).GetNoDataValue()
    min_value = 0
    max_value = 0
    if max_min == True:
        array_band = ds.GetRasterBand(1).ReadAsArray(0, 0, ds_width, ds_height).astype(np.float64)
        array_band_masked = ma.masked_values(array_band, ds_nodata)  # nodata值不参与运算
        min_value = np.min(array_band_masked)  # 获取最小值
        max_value = np.max(array_band_masked)   # 获取最大值
    return ds_width, ds_height, ds_bands, ds_nodata, min_value, max_value
    # print("影像的宽度为：" + str(ds_width))
    # print("影像的高度为：" + str(ds_height))
    # print("仿射地理变换参数为：" + str(ds_geo))
    # print("投影坐标系为：" + str(ds_prj))
    # data = ds.ReadAsArray(0, 0, ds_width, ds_height)  # 以数组的形式读取整个数据集


def Normalization_value(filepath, out_path, obverse=True):
    """
    归一化处理函数

    参数：
    filepath(str): 输入影像的路径
    out_path(str): 输出影像的路径
    obverse(bool, optional): 是否进行正向归一化，默认为True

    返回值：
    无

    """

    # ------------------------------------影像基础信息------------------------------------
    ds = gdal.Open(filepath)  # 打开数据集dataset
    ds_width = ds.RasterXSize  # 获取数据宽度
    ds_height = ds.RasterYSize  # 获取数据高度
    ds_geo = ds.GetGeoTransform()  # 获取仿射地理变换参数
    ds_prj = ds.GetProjection()  # 获取投影信息
    nodata=ds.GetRasterBand(1).GetNoDataValue()  # 获取影像的nodata值

    driver = gdal.GetDriverByName('GTiff')  # 载入数据驱动，用于存储内存中的数组
    ds_result = driver.Create(out_path, ds_width, ds_height, bands=ds.RasterCount, eType=gdal.GDT_Float64)
    if ds.RasterCount == 1: #如果是单波段影像
        array_band = ds.GetRasterBand(1).ReadAsArray(0, 0, ds_width, ds_height).astype(np.float64)
        array_band_masked = ma.masked_values(array_band, nodata) #nodata值不参与运算
        min_value = np.min(array_band_masked) #ds.GetRasterBand(1).GetMinimum()   # 获取最小值
        max_value = np.max(array_band_masked) #ds.GetRasterBand(1).GetMaximum()   # 获取最大值
        mask=np.where(array_band == nodata, 0, 1) # 设置一个蒙版
        if obverse==True: #正向标准化
            array_result = mask * (array_band_masked - min_value) / (max_value - min_value) #不用mask的形式，最大值最小值计算会出错，好像是溢出了
        else: #反向标准化
            array_result = mask * (max_value - array_band_masked) / (max_value - min_value)
        #array_result.mask = ma.nomask #恢复被mask的数字，这句话好像也不用
        array_result = nodata * (1 - mask) + array_result  # 恢复图像中的nodata
        ds_result.GetRasterBand(1).WriteArray(array_result)  # 将结果写入数组
        array_result = None
        del array_result
    else: #如果是多波段影像
        print("正在遍历所有波段......")
        for i in range(1, ds.RasterCount):  # 遍历所有波段
            array_band = ds.GetRasterBand(i).ReadAsArray(0, 0, ds_width, ds_height).astype(np.float64)
            array_band_masked = ma.masked_values(array_band, nodata)  # nodata值不参与运算
            min_value = np.min(array_band)  # 获取最小值
            max_value = np.max(array_band)  # 获取最大值
            if obverse == False:  # 正向标准化
                array_result = (array_band_masked - min_value) / (max_value - min_value)
            else:  # 反向标准化
                array_result = (max_value - array_band_masked) / (max_value - min_value)
            #array_result.mask = ma.nomask #恢复被mask的数字
            ds_result.GetRasterBand(i).WriteArray(array_result)  # 将结果写入数组
            array_result = None
            del array_result
    ds_result.SetGeoTransform(ds_geo)  # 导入仿射地理变换参数
    ds_result.SetProjection(ds_prj)  # 导入投影信息
    ds_result.GetRasterBand(1).SetNoDataValue(nodata)  # 设置nodata值
    del ds_result # 删除内存中的结果，否则结果不会写入图像中
    # 删除内存中的结果，否则结果不会写入图像中
    print("归一化处理已完成......")

def Normalization_ndarray(data_array: np.ndarray, begin=1, nodata=-3.40282346639e+38):
    # 归一化处理函数，处理传入的numpy数组
    layer_count = data_array.shape[0]  # 获得图层数量，第一维 # len(data_array)

    for i in range(begin, layer_count): # begin 因为我将label（target）和data（layer）写到了同一个ndarray中，所有要从1开始，第0个为target，它不能为离散型变量
        masked_array = ma.masked_values(data_array[i,:,:], nodata) #nodata值不参与运算
        min_value = np.min(masked_array)  # 获取最小值
        max_value = np.max(masked_array)  # 获取最大值
        mask = np.where(data_array[i,:,:] == nodata, 0, 1) # 设置一个蒙版
        # 全部正向标准化
        array_result = mask * (masked_array - min_value) / (max_value - min_value)
        #不用mask的形式，最大值最小值计算会出错，好像是溢出了
        data_array[i,:,:] = nodata * (1 - mask) + array_result  # 恢复图像中的nodata

    return data_array

    print("numpy数组归一化处理已完成......")

if __name__ == "__main__":
    print("\n")
    # os.environ['PROJ_LIB'] = 'G:/Anaconda/envs/pyDL/Lib/site-packages/osgeo/data/proj'
    # os.environ['GDAL_DATA'] = 'G:/Anaconda/envs/pyDL/Lib/site-packages/osgeo/data'
    # 添加PROJ至环境变量，消除警告
    file_path = r"D:\test\预测分析\距离水系.tif"  # 输入的栅格数据路径
    out_path1 = r"D:\test\预测分析\距离水系_norm.tif"  # 导出的文件路径
    Get_data(file_path)  # 执行函数，获取影像基本信息，可以注释掉
    Normalization_value(file_path, out_path1,False)
    print("--------------------------------------程序结束--------------------------------------")